from torch.utils.data import Dataset
import os
from PIL import Image
import torch
from torchvision.transforms import Compose,Resize,ToTensor,Normalize
from utils import IMAGESUFFIX

class Market1501(Dataset):
    def __init__(self,
                 begin_pid = 0,
                 pre_process=None,
                 data_dir = None,
                 extra_annotations =None,
                 mode = 'train',
                 color_space = 'RGB',
                 ):
        super().__init__()
        assert mode in ['test','train']
        self.mode = mode
        self.color_space = color_space
        self.begin_pid = begin_pid
        
        query_dir   = os.path.join(data_dir,'query')
        gallery_dir = os.path.join(data_dir,'bounding_box_test')
        train_dir   = os.path.join(data_dir,'bounding_box_train')


        person_id_container = set()

        _dirs = [query_dir,gallery_dir] if mode == 'test' else [train_dir]

        for directory in _dirs:
            for file in os.listdir(directory):
                if file.endswith(IMAGESUFFIX):
                    features = file.split('_')
                    person_id = int(features[0])
                    if person_id == -1:
                        continue
                    person_id_container.add(person_id)

        self.person_id2label = {pid:self.begin_pid+label for label,pid in enumerate(sorted(person_id_container))}

        if self.mode == 'test':
            self.query_data = self.get_data(query_dir)
            self.gallery_data = self.get_data(gallery_dir)

            self.data = self.query_data + self.gallery_data

        elif self.mode == 'train':
            self.data = self.get_data(train_dir)

        if pre_process is None:
            self.pre_process = Compose([
                                        Resize((256,128)),
                                        ToTensor(),
                                        Normalize(mean=(0.485,0.456,0.406),
                                                std=(0.229,0.224,0.225))
                                    ])
        else:
            self.pre_process = pre_process
    
    def get_data(self,dir):
        data = []
        for file in os.listdir(dir):
            if file.endswith(IMAGESUFFIX):
                features = file.split('_')
                person_id = int(features[0])
                if person_id == -1:
                    continue
                assert len(features[1]) == 4
                camera_id = int(features[1][:2][1:])
                video_id = int(features[1][2:][1:])

                assert 0 <= person_id <= 1501
                assert 1 <= camera_id <= 6
                
                data.append({
                    'image_path': os.path.join(dir, file),
                    'person_id': self.person_id2label[person_id],
                    'camera_id': camera_id,
                    'video_id': video_id,
                    'original_pid': person_id
                })
        return data
    
    def __len__(self):
        return len(self.data)
    def __getitem__(self, index):
        data = self.data[index]
        image_path = data['image_path']
        person_id  = data['person_id']
        camera_id  = data['camera_id']
        video_id   = data['video_id']

        image = Image.open(image_path).convert(self.color_space)
        
        image = self.pre_process(image)
        
        return image,person_id,camera_id,video_id,image_path